import chromadb

client = chromadb.Client()
#chroma_client = chromadb.PersistentClient(path="./db")
# Create a new chroma collection
collection_name = "filter_example_collection"
collection = client.create_collection(name=collection_name)

# Add some data to the collection
collection.add(
    embeddings=[
        [1.1, 2.3, 3.2],
        [4.5, 6.9, 4.4],
        [1.1, 2.3, 3.2],
        [4.5, 6.9, 4.4],
        [1.1, 2.3, 3.2],
        [4.5, 6.9, 4.4],
        [1.1, 2.3, 3.2],
        [4.5, 6.9, 4.4],
    ],
    metadatas=[
        {"status": "read"},
        {"status": "unread"},
        {"status": "read"},
        {"status": "unread"},
        {"status": "read"},
        {"status": "unread"},
        {"status": "read"},
        {"status": "unread"},
    ],
    documents=["A document that discusses domestic policy", "A document that discusses international affairs", "A document that discusses kittens", "A document that discusses dogs", "A document that discusses chocolate", "A document that is sixth that discusses government", "A document that discusses international affairs", "A document that discusses global affairs"],
    ids=["id1", "id2", "id3", "id4", "id5", "id6", "id7", "id8"],
)

# Get documents that are read and about affairs
res1 = collection.get(where={"status": "read"}, where_document={"$contains": "affairs"})
print(res1)
# Get documents that are about global affairs or domestic policy
res2 = collection.get(where_document={"$or": [{"$contains": "global affairs"}, {"$contains": "domestic policy"}]})
print(res2)

# Get 5 closest vectors to [0, 0, 0] that are about affairs
# Outputs 3 docs because collection only has 3 docs about affairs
res3 = collection.query(query_embeddings=[[0, 0, 0]], where_document={"$contains": "affairs"}, n_results=5)
print(res3)

res4 = collection.query(query_embeddings=[[0, 0, 0]], where_document={"$not_contains": "domestic policy"}, n_results=5)
print(res4)




